import cv2
import numpy as np
import scipy.fftpack as fftpack


# 傅里叶变换（FFT）
def fft_transform(image):
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = 20 * np.log(np.abs(fshift))
    return magnitude_spectrum


# 离散余弦变换（DCT）
def dct_transform(image):
    dct = cv2.dct(np.float32(image))
    return dct


# 理想低通滤波
def ideal_low_pass_filter(image, cutoff):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    mask = np.zeros((rows, cols), np.uint8)
    mask[crow - cutoff:crow + cutoff, ccol - cutoff:ccol + cutoff] = 1
    fshift = fshift * mask
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back


# 巴特沃斯低通滤波
def butterworth_low_pass_filter(image, cutoff, order):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    u, v = np.meshgrid(np.arange(cols), np.arange(rows))
    d = np.sqrt((u - ccol) ** 2 + (v - crow) ** 2)
    h = 1 / (1 + (d / cutoff) ** (2 * order))
    fshift = fshift * h
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back


# 理想高通滤波
def ideal_high_pass_filter(image, cutoff):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    mask = np.ones((rows, cols), np.uint8)
    mask[crow - cutoff:crow + cutoff, ccol - cutoff:ccol + cutoff] = 0
    fshift = fshift * mask
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back


# 拉普拉斯频域滤波
def laplacian_frequency_filter(image):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    u, v = np.meshgrid(np.arange(cols), np.arange(rows))
    h = -(u - ccol) ** 2 - (v - crow) ** 2
    fshift = fshift * h
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back


# 带通滤波
def band_pass_filter(image, low_cutoff, high_cutoff):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    mask = np.zeros((rows, cols), np.uint8)
    u, v = np.meshgrid(np.arange(cols), np.arange(rows))
    d = np.sqrt((u - ccol) ** 2 + (v - crow) ** 2)
    mask[(d >= low_cutoff) & (d <= high_cutoff)] = 1
    fshift = fshift * mask
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back


# 带阻滤波
def band_stop_filter(image, low_cutoff, high_cutoff):
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    mask = np.ones((rows, cols), np.uint8)
    u, v = np.meshgrid(np.arange(cols), np.arange(rows))
    d = np.sqrt((u - ccol) ** 2 + (v - crow) ** 2)
    mask[(d >= low_cutoff) & (d <= high_cutoff)] = 0
    fshift = fshift * mask
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)
    return img_back